[Spark 4] Streaming runner proof of concept, unbounded sources, windowed GBK and stateful ParDo on transformWithState - #39576
Draft
tkaymak wants to merge 10 commits into
Draft
[Spark 4] Streaming runner proof of concept, unbounded sources, windowed GBK and stateful ParDo on transformWithState#39576tkaymak wants to merge 10 commits into
tkaymak wants to merge 10 commits into
Conversation
Prepares the structured streaming runner for a Spark 4 streaming translator without changing batch behavior. - Add PipelineTranslatorFactory, the seam the Spark 4 module shadows to dispatch streaming pipelines. The shared base rejects streaming with a clear message instead of the previous generic checkArgument. - Open up EvaluationContext (non final, protected ctor, leaves()) and add a no-op stop() so a streaming context can override evaluation. - Add a createEvaluationContext hook to PipelineTranslator and skip the persist and lineage breaking optimizations for streaming datasets. - Plumb the EvaluationContext into SparkStructuredStreamingPipelineResult so cancel() can stop a running streaming query. - Default the state store provider to RocksDB, required by Spark 4 transformWithState and inert for batch. - Add watermarkDelayMillis, maxRecordsPerMicroBatch, maxBatchDurationMillis and streamingStopAfterIdleBatches options. - Narrow the Spark 4 test source override exclude to the legacy DStream package. The previous glob also matched structuredstreaming and would have silently dropped the structured streaming tests.
Wraps any Beam UnboundedSource as a Spark 4 DataSourceV2 micro-batch streaming source. The rows have exactly two columns, payload of type BINARY holding the element encoded with a FullWindowedValueCoder, and eventTimestamp of type TIMESTAMP holding the Beam event timestamp, so no Catalyst encoder has to be generated for Beam types. Offsets are opaque epoch counters. The driver never reads from the source and never inspects its progress, latestOffset simply returns the previous epoch plus one so Spark keeps planning micro-batches. Termination is therefore owned by the lifecycle, not by the offsets. The source is split once on the driver and the sub sources are memoized, so every micro-batch plans the same stable set of partitions. On the executor a static cache keyed by checkpoint location, source id and split id keeps the Beam UnboundedReader alive across micro-batches, mirroring the legacy MicrobatchSource. Checkpoint marks are cached in executor memory only and are never committed durably, which is documented POC scope with weak failure recovery. UnboundedSourceDataset is the translator facing entry point. It applies withWatermark exactly once, since Spark 4 forbids a second declaration further down the plan. The test confirms that the EventTimeWatermark logical node survives one and two typed maps in both the logical and the analyzed plan, and that the running query really tracks the watermark afterwards. It also reads a finite set of elements end to end and checks the decoded payloads and timestamps.
Adds a single generic transformWithState operator that can host any keyed Beam transform, both the stateful ParDo stack and the group also by window stack that implements a windowed GroupByKey. The mode is selected by BeamStatefulProcessorConfig, everything else is shared. Keys, inputs and outputs are raw Beam coder bytes and every Spark encoder involved is BINARY or STRING, so Catalyst never has to derive a schema for a Beam type. TwsTransformFactory owns the row layout and documents it. Beam state is bridged by TwsStateInternals on top of BytesKV, a tiny string to bytes store backed by exactly one Spark MapState. One map is a requirement rather than a preference, ReduceFnRunner invents state tags at runtime so the set of state addresses is not known when init has to declare its state variables. Timers are bridged by TwsTimerInternals, which keeps the full TimerData in a second map and registers only bare wake ups with Spark, reconciling them against listTimers so an expiry can never be registered twice. Two boundary conditions needed care. Spark deletes the expiry it is firing after the callback returns, so a wake up re armed at that same millisecond is nudged one millisecond forward. Spark also expires a wake up as soon as the expiry is at or before the batch watermark, while Beam only fires an event time timer once the watermark is strictly past it, so timers sitting exactly on the watermark are withheld and re armed instead of being handed to Beam, which would decline to fire them and silently drop the on time pane. Processing time timers are rejected with a clear message, transformWithState runs in a single TimeMode and this operator uses event time. The state and timer bridges are unit tested against an in memory store, and BeamStatefulProcessorTest runs both modes through a real streaming query on a live SparkSession.
Shadows PipelineTranslatorFactory so streaming pipelines dispatch to a new PipelineTranslatorStreaming, which extends PipelineTranslatorBatch to reuse Impulse, Window.Assign, Flatten, Reshuffle, the bounded read and stateless ParDo unchanged, while registering placeholders for the unbounded read, GroupByKey and stateful ParDo that WS-D2 will replace. Combine.PerKey is deliberately left unregistered so it auto-expands into GroupByKey plus ParDo. Adds StreamingEvaluationContext, which starts one noop-sink streaming query per leaf dataset with a checkpoint directory derived from the pipeline options, registers an idle-stop listener for test termination, and makes stop() idempotent and safe to call concurrently with evaluate() so cancel() can interrupt a running streaming pipeline from another thread.
…etons Adds StreamingTestUtils, the ListBackedUnboundedSource plus static collector and pipeline option factories the streaming translators will be tested against, and five skeleton test classes for the stateless ParDo, windowed GroupByKey, stateful ParDo, chained stateful, and lifecycle scenarios. Every skeleton compiles and runs, reporting as skipped: each is Ignore'd with a comment naming the exact WS-D2 translator it needs once that lands.
Replaces the WS-D2 placeholders in PipelineTranslatorStreaming with real translators for the unbounded read, GroupByKey and stateful ParDo. ReadUnboundedTranslator wraps the Beam UnboundedSource with UnboundedSourceDataset, which already declares the single withWatermark for the whole query, and decodes the binary payload column into the Dataset<WindowedValue<T>> shape every other translator consumes. The event timestamp column is projected away, the EventTimeWatermark plan node survives below the projection and transformWithState reads the query wide watermark rather than a column. GroupByKeyStreamingTranslator and StatefulParDoStreamingTranslator both encode their input into the byte[] row layout of TwsTransformFactory, hand it to the generic transformWithState operator in GROUP_ALSO_BY_WINDOW and STATEFUL_PARDO mode respectively, and decode the tagged output rows back out. The stateful ParDo splits the output by tag index and skips additional outputs that nothing consumes, so an unused tag does not start a second streaming query. StreamingTranslationHelpers holds the shared guards and row conversions. Pipelines outside the POC scope are rejected at translation time with a message naming the feature, specifically merging windows, custom triggers, accumulating panes, processing time timers, non deterministic key coders, non KV input to a stateful ParDo, side inputs on a stateful ParDo, @OnWindowExpiration and @RequiresTimeSortedInput.
Turns the five Ignore'd streaming test skeletons into real, asserting tests against the WS-D2 translators. Every skeleton declared its pipeline as a local TestPipeline variable, which TestPipeline itself rejects at run() because it is not a Rule field. These tests cannot use PAssert anyway, a streaming pipeline here never gets a final watermark so panes never finalize, and each test needs its own options, so TestPipeline buys nothing. They now build a plain Pipeline.create(options) and assert after waitUntilFinish against the static collector in StreamingTestUtils. The unused TestPipeline factory is dropped from StreamingTestUtils and the reasoning is recorded in its javadoc. StatelessParDoStreamingTest and StreamingPipelineLifecycleTest gain the same relaxed Kryo SparkSessionRule as the rest of the suite, so that the whole package shares one session and a cancelled or idle stopped query cannot take that session down with it. Asserted results, all derived from each test's own input list. Stateless ParDo, exactly 0, 2, 4 up to 18. Fixed 10s window count per key, a=2 and b=1. Sliding 10s every 5s, two a=2 panes, one per shared window. Late data dropped, a=1 and sentinel=1, never a=2. Stateful ParDo, a, b, c once each plus two timer sentinels. Chained dedup into windowed sum, a=8 and b=10. The chained case is the one that matters, a=13 would mean the dedup state was lost and an empty result would mean the watermark got stuck between the two operators, so exactly a=8 and b=10 is the observable proof of cross operator watermark propagation. The late data test is the only one that depends on how the source round robins elements across splits, so it pins maxRecordsPerMicroBatch to 1, spells out the resulting per batch schedule in a comment, and asserts the split count assumption up front rather than leaving it silent. The lifecycle tests wait for a query to actually be active before cancelling, since translation is asynchronous and an early cancel would find a null evaluation context, and they check the query really stopped rather than only that the state was relabelled.
Spark 4 broadcasts a StateSchemaMetadata to the executors for every transformWithState query, using the user Kryo instance. Nothing registered that class, so any Beam streaming pipeline with state or timers died on its first micro-batch as soon as spark.kryo.registrationRequired was on, which is the default for this module's tests and a perfectly reasonable production setting. The same applies to MemoryWriterCommitMessage, the commit message of Spark's memory sink, which is nested inside the already registered DataWritingSparkTaskResult. Both are registered by name, because the shared runner base also compiles against Spark 3 where neither class exists, and both are registered with a JavaSerializer rather than Kryo's default field serializer. Both are Scala case classes holding further Scala and Spark types, immutable.Map, StructType, avro Schema and Row, none of which are registered either, so going through Java serialization covers the whole object graph at once instead of forcing this list to track Spark's internal field layout. Neither object is on a hot path. Every streaming test now runs with the module default of spark.kryo.registrationRequired=true, the per test relaxations are gone. SparkKryoRegistratorStreamingTest names both classes at compile time against the Spark 4 classpath, so a future rename becomes a compile error rather than a silently dropped registration.
Two javadoc paragraphs in the shared base drifted from google-java-format, one from the streaming note added to SparkStructuredStreamingRunner and one from the new Kryo registrations. The shared base sources under runners/spark/src are checked by the :runners:spark project rather than by :runners:spark:3 or :runners:spark:4, whose own spotless targets only see their per version override directories, so it is easy to miss.
ChainedStatefulStreamingTest asserts what the chained pipeline computes. It cannot distinguish a single query holding two transformWithState operators from two queries that happen to add up to the same numbers, and it cannot show whether the watermark moved at all or whether the whole input simply landed in one micro-batch where everything is trivially on time. ChainedStatefulStreamingEvidenceTest asserts the run itself, against Spark's own StreamingQueryProgress. It pins one record per split per micro-batch so the watermark has to climb in steps, records every progress event through a listener, and then asserts that all of them carry one query id, that one micro-batch reports exactly two transformWithStateExec state operators, and that the watermark takes at least three strictly increasing values. The second test adds a tap between the two operators, so the late record is observed leaving the dedup operator and absent from the windowed sum. That is the difference between a record excluded for lateness and a record that never arrived. Both tests print the raw per batch progress they recorded, since that output is the evidence the phase gate report quotes. Also drops four javadoc references to a Kryo relaxation that no longer exists anywhere in the suite.
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Draft, opened for design discussion rather than for merging as one lump. See "How I would like to split this" at the bottom.
Addresses #36841.
What this is
The Spark 4 runner merged in #38255 is batch only.
SparkStructuredStreamingRunnerrejects streaming outright. This branch is a proof of concept that Spark 4'stransformWithStatecan host the Beam streaming model, and it turns out it can, including several chained stateful operators inside a single Structured Streaming query with a watermark that propagates correctly across them.I would like to agree on the approach before splitting this into reviewable pieces.
Approach
Dispatch seam. The shared base gains
translation/PipelineTranslatorFactory.create(boolean streaming), which throws for streaming and names the Spark 4 module.runners/spark/4/srcshadows that one file through the existing source override copy and returns a streaming translator instead. Every piece oftransformWithState,StreamingQueryand DataSourceV2 streaming code lives only underrunners/spark/4/src, because the shared base still compiles against Spark 3.5 wheretransformWithStatedoes not exist. All shared base edits are behaviour preserving for Spark 3.Source. A DataSourceV2
TableProviderwrapping anyUnboundedSource. Offsets are opaque epoch counters, the driver never reads them. Executor sidePartitionReaders holdUnboundedReaders in a static cache and resume from a locally cachedCheckpointMark, in the style of the legacyMicrobatchSource. The row schema ispayload BINARYpluseventTimestamp TIMESTAMP, so no new Catalyst encoder work was needed.Watermark. Declared exactly once, in the read translator, on the raw rows before the typed decode. Spark forbids re-declaring it. The
EventTimeWatermarkplan node survives the projection and the typed map, which was the main thing I was unsure about going in. There is a test that asserts this at both the plan level and at runtime.Stateful execution. One generic
transformWithStateoperator,StatefulProcessor<byte[], byte[], byte[]>withEncoders.BINARY()throughout, hosting anyDoFnthroughDoFnRunners. Stateful ParDo goes throughsimpleRunnerplusdefaultStatefulDoFnRunner, GBK throughGroupAlsoByWindowViaWindowSetNewDoFnplusSystemReduceFn.bufferingandKeyedWorkItems, which follows the FlinkWindowDoFnOperatorrecipe. State is a port of the legacySparkStateInternalsonto a singleMapState<String, byte[]>keyed by namespace plus tag. That layout is the only one I found that can hostReduceFnRunner's dynamically created system tags. Per@StateIdcolumn families would be faster and are an obvious follow up.Lifecycle.
StreamingEvaluationContextstarts one query per leaf against thenoopsink and blocks until all are terminal.cancel()reaches it through anAtomicReferenceset by the async translate task. Tests terminate through an idle stop listener that gracefully stops a query after N consecutive empty micro batches.Results
:runners:spark:3:test:runners:spark:4:test121 of the Spark 4 tests are structured streaming. The streaming suite was run four separate times with identical results.
Asserted end to end: stateless ParDo, fixed and sliding window GBK, late data dropping, stateful dedup with an event time timer, chained stateful then windowed sum, and pipeline lifecycle through RUNNING, DONE and CANCELLED.
The chained case is the one that matters, and it is asserted against Spark's own
StreamingQueryProgressrather than only against the computed values: one query id across every batch, exactly twotransformWithStateExecoperators per batch, and four strictly increasing watermarks. That rules out two separate queries that happen to sum correctly, and it rules out everything landing in one batch where all data is trivially on time.One finding worth flagging on its own
Spark expires a
transformWithStatewake up atexpiry <= batchWatermark. Beam fires an event time timer only once the watermark is strictly past it, andAfterWatermark.pastEndOfWindowis strict as well. Handing Beam a timer one millisecond early is therefore not merely early, it is destructive: the trigger declines, the runner is entitled to assume the timer will not be redelivered, and the on time pane is silently lost with no error anywhere. The end of window timer of a fixed window sits at exactlywindow.maxTimestamp(), so this hits most windows.Fixed by bounding the fire set at
min(firedExpiry, watermark - 1)and re-arming the withheld timer atfiredExpiry + 1. Anyone else bridging Beam timers ontotransformWithStatewill hit this.Rejected at translation time, with a named message
Merging and session windows, custom triggers including
AfterWatermarkwith early or late firings, accumulating panes, processing time timers,@OnWindowExpiration,@RequiresTimeSortedInput, side inputs on a stateful ParDo, non deterministic key coders.Known limitations
PAsserton the streaming path. No source emits a final positive infinity watermark, so panes never finalise. Streaming tests assert against a static collector afterwaitUntilFinish. I would like input on what reviewers expect here, since this is the biggest departure from how the rest of Beam tests runners.stop()does not drain. An in flight micro batch finishes, anything not yet pulled is left unprocessed.Combine.PerKey.Two smaller things in here
SparkSessionFactorynow registersStateSchemaMetadataandMemoryWriterCommitMessagewith Kryo, by name since the shared base also compiles against Spark 3. Spark 4 broadcasts aStateSchemaMetadatafor everytransformWithStatequery, so without this any stateful streaming pipeline dies on its first micro batch underspark.kryo.registrationRequired=true.runners/spark/4/build.gradlenarrowsexclude "**/translation/streaming/**"toexclude "**/runners/spark/translation/streaming/**". The old glob also matchedstructuredstreaming/translation/streaming, so it would have silently dropped every new test in this branch from the Spark 4 module.Not done yet, deliberately
@Ignore("TODO: Reactivate with streaming."). They cannot be un-ignored in place, since the shared base test tree is compiled by both modules and Spark 3 still correctly throws. Reactivating them means moving them into the Spark 4 override tree. Happy to do that, it just needs a decision on where they should live.How I would like to split this
Assuming the approach is acceptable, five PRs, each independently green:
Slices 4 and 5 would benefit from a reviewer who knows Spark's Structured Streaming internals, not only Beam.